__author__ = 'tony_su'

import tempfile
import datetime
import hashlib
import argparse
import sys
import subprocess
import logging
import json
import traceback
import dd_legacy_dbtool
import misc_utils
import psycopg2
from librule import Type, LibruleWrapper
from common_dl import *



class deylistException(Exception):pass

class DenyList(object):
    def __init__(self):

        _p = misc_utils.get_db_params()
        self._db_con = dd_legacy_dbtool.DBConnection(_p['database'], _p['user'], _password='', _app_name='custdenylist')
        self.type_supported = [Type.IP,Type.URL,Type.SHA1,Type.DOMAIN]

        self.DB_INSERT = 'insert'
        self.DB_DELETE ='delete'
        self.DB_QUERY ='query'
        self.DB_UPDATE='update'

        self.librule = LibruleWrapper()

    def check_type(self,type):
        try:
            type = int(type)
        except:
            return None
        else:
            if type in self.type_supported:
                return type
        return None

    def create_denylist_object(self,type=None,data=None):
        denylist={}
        for i in self.type_supported:
            denylist[i]=[]

        if type is not None:
            denylist[type].append(data)

        return denylist

    def _exec_usbxcli_custdenylist(self, ini_file=None):
        '''
        Execute a usbxcli set-custdenylist command
        :param: denylist file
        :return:
        '''
        cmd = []
        cmd.append('/usr/bin/python')
        cmd.append('/var/app_data/U-Sandbox/u-sandbox/usandbox/cli/usbxcli.py')
        cmd.append('set-custdenylist')

        if ini_file:
            cmd.append(ini_file)

        proc, stdout, stderr = (None,) * 3

        logging.info("cmd=[%s]" % (" ".join(map(str, cmd))))
        try:
            proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            (stdout, stderr) = proc.communicate()
        except:
            logging.error('Fail to run command  {}'.format(cmd))
            logging.debug(traceback.format_exc())
            return {'status': False, 'msg': stderr, 'retcode': -9999}

        return {'status': True, 'msg': stdout, 'retcode': proc.returncode}

    def get_usbx_invalid_denylist(self):

        logging.debug('Check InvalidCustDenyList')

        result = self._exec_usbxcli_custdenylist()
        try:
            invalid_data = json.loads(result['msg'])['InvalidCustDenyList']
        except Exception as e:
            logging.error('Get U-sandbox denylist error {}'.format(e))
            return None

        map={
            'Domain':Type.DOMAIN,
            'IP':Type.IP,
            'SHA1':Type.SHA1,
            'URL':Type.URL
        }

        data = {}
        for t, v in invalid_data.iteritems():
            data[map[t]] = v

        return data

    def _loop_list(self,f,items):

        if not isinstance(items,list):
            raise deylistException('This {} is not an python list object'.format(items))

        for i in items:
            # logging.debug('Write {} in temp denylist ini'.format(i))
            f.write('{}{}'.format(i, '\n'))

    def write_ini(self,denylist,tfname):
        '''
        a dict object
        :param denylist:
        :return:
        '''
        if not isinstance(denylist,dict):
            raise deylistException('This {} is not an python dict object'.format(denylist))

        header = '[Main]\nVersion = 1\nLastUpdate = {}\n'
        with open(tfname, 'a') as f:
            f.write(header.format(datetime.datetime.now().strftime("%Y-%m-%dT%H:%M:%SZ")))
            f.write('[Rules IP]\n')
            if Type.IP in denylist:
                self._loop_list(f,denylist[Type.IP])
            f.write('[Rules DNS]\n')
            if Type.DOMAIN in denylist:
                self._loop_list(f,denylist[Type.DOMAIN])
            f.write('[Rules URL]\n')
            if Type.URL in denylist:
                self._loop_list(f,denylist[Type.URL])
            f.write('[Rules SHA1]\n')
            if Type.SHA1 in denylist:
                self._loop_list(f,denylist[Type.SHA1])

    def set_usbx_denylist(self,ini):

        logging.debug('Import data to custdenylist')

        result = self._exec_usbxcli_custdenylist(ini)['msg'].splitlines()
        status = False
        total, fail =0, 0
        for i in result:
            if i.startswith('CustDenyList imported successfully.'):
                status = True
            if i.startswith('Total:'):
                total = int(i.split(':')[1])
            if i.startswith('Failure:'):
                fail = int(i.split(':')[1])
        if not status:
            logging.error('Call usbxcli error')
            raise deylistException(result[0])
        return total,fail

    def db_insert_try(self,denylist):
        # commit will be done in outer side
        for type,v in denylist.iteritems():
            for data in v:
                self._db_con.execute('''INSERT INTO tbl_deny_list (data,type,timestamp) VALUES(%s,%s,%s);''',
                                    (data, type, datetime.datetime.utcnow()))

    def db_update_try(self,data,id):
        # commit will be done in outer side

        self._db_con.execute('UPDATE tbl_deny_list set data=%s where id=%s',(data,id))
        if self._db_con.rowcount() :
            return True
        return False

    def db_get(self):
        self._db_con.execute('SELECT data,type from tbl_deny_list ORDER BY type;')
        result = self._db_con.fetchall()

        denylist = self.create_denylist_object()
        for data, type in result:
            denylist[type].append(data)
        return denylist

    def db_get_cnt(self):
        self._db_con.execute('SELECT count(id) from tbl_deny_list;')
        result = self._db_con.fetchone()
        if result:
            return int(result[0])
        return 0

    def db_del_by_data(self, data):
        try:
            for t, v in data.iteritems():
                for i in v:
                    logging.error('delete invalid items by data {} {}'.format(i, t))
                    self._db_con.execute('''DELETE FROM tbl_deny_list where data=%s''', (i,))
            self._db_con.commit()
        except psycopg2.Error:
            logging.error(traceback.format_exc())
            self._db_con.rollback()
            return False
        logging.debug('Success to update denylist table')
        return True

    def db_del_by_id_try(self, ids):
        '''
        Try to delete deny list by id
        or trancate the tbl_deny_list table if ids is empty
        :param ids:
        :return:
        '''
        # commit will be done in outer side

        is_affect = False

        if len(ids) == 0:
            logging.debug('Try to delete all denylist')
            self._db_con.execute('''TRUNCATE tbl_deny_list;''')
            is_affect = True
        else:
            logging.debug('Try to delete id {}'.format(ids))
            for i in ids.split(','):
                logging.error('delete items by id {}'.format(i))
                self._db_con.execute('''DELETE FROM tbl_deny_list where id=%s''', (i,))
                if not is_affect and self._db_con.rowcount():
                    is_affect = True
        return is_affect

    def db_update_timestamp(self,timestamp):
        if not timestamp:
            logging.debug('Do not update timestamp for sync')
            return True
        try:
            self._db_con.execute("update  tb_global_setting set value = (extract(epoch from now()))::bigint where key = 'configuration.denylist.timestamp';")
            self._db_con.commit()
        except psycopg2.Error:
            self._db_con.rollback()
            logging.error('Fail to update timestamp {}'.format(traceback.format_exc()))
            return False
        logging.debug('Success to update timestamp')
        return True

    def sort_denylist(self,data):

        denylist = self.create_denylist_object()

        for i in data:
            rule_type = int(i['type'])
            # logging.debug('find item {}'.format(i))
            if rule_type in self.type_supported:
                denylist[rule_type].append(i['data'])
            else:
                logging.error('Invalid type {}'.format(i['type']))
                return None
        return denylist

    def create_dumy_denylist_data(self,cnt=20):
        '''
        {
            { "data":"1.1.1.1","type":1},
            { "data":"1.1.1.2","type":1}
        }

        :param cnt:
        :return:
        '''
        denylist_data=[]

        try:
            cnt = int(cnt)
        except:
            cnt= 20
        for i in xrange(cnt):
            items={}
            items['type']= i%3
            if items['type'] == Type.URL:
                items['data']='http://baidu{}.com'.format(i)
            elif items['type'] == Type.SHA1:
                items['data'] = hashlib.sha1(str(i)).hexdigest()
            elif items['type'] == Type.DOMAIN:
                items['data'] ='baidu{}.com'.format(i)
            elif items['type'] == Type.IP:
                items['data'] = 'baidu{}.com'.format(i)
            else:
                print 'Invalide type'
            denylist_data.append(items)

        return json.dumps(denylist_data)

    def is_duplicate_db(self,data,type):

        denylist = self.db_get()
        self.librule.import_rules(denylist)
        return self.librule.match_all(data,type)

    def _sync_denylist(self,db_action,op_data = None,timestamp=True):
        '''
        Place db operation and usbx operation to one transaction
        1. Save deny list to db/usbx
        2. or sync db deny list to usbx
        :param db_action:
            insert: insert deny list to db
            delete: delete deny list by id
        :param op_data
        :return:
        '''
        logging.debug('db_action:[{}], op_data:[{}]'.format(db_action,op_data))
        try:
            if db_action== self.DB_INSERT:
                # insert data to tbl_deny_list
                logging.debug('Try to save denylist to db')
                self.db_insert_try(op_data)
                if self.db_get_cnt() > MAX_DENYLIST_CNT:
                    logging.error('The denylist exceed the max count [{}]'.format(MAX_DENYLIST_CNT))
                    raise codeException(Code.TOO_MANY_ITEM)
            elif db_action == self.DB_DELETE:
                if not self.db_del_by_id_try(op_data):
                    logging.warn('There no data to be deleted')
                    raise codeException(Code.DB_NO_DATA_UPDATE_DELETE)
            elif db_action == self.DB_UPDATE:
                data,id = op_data
                if not self.db_update_try(data,id):
                    logging.warn('There is not data to be updated')
                    raise codeException(Code.DB_NO_DATA_UPDATE_DELETE)

            # get deny list from tbl_deny_list for sync operation
            if db_action in [self.DB_DELETE,self.DB_QUERY,self.DB_UPDATE]:
               logging.debug('Get data from tbl_deny_list')
               op_data = self.db_get()

            logging.debug('Try to sync usbx denylist')
            tf = tempfile.NamedTemporaryFile(dir='/tmp',prefix='cusdenylist_')
            self.write_ini(op_data,tf.name)
            total, fail = self.set_usbx_denylist(tf.name)

        except codeException as e:
            self._db_con.rollback()
            logging.warn('Encounter code [{}]'.format(e.code))
            return e.code
        except:
            self._db_con.rollback()
            logging.error('Fail to save denylist to db and usbx [{}]'.format(traceback.format_exc()))
            return Code.SAVE_DENY_LIST_ERROR

        self._db_con.commit()
        logging.debug('Success to save/sync {} denylist to tbl_deny_list/usbx'.format(total))

        if fail != 0:
            logging.error('Find invalid items [{}]'.format(fail))
            self.check_usbx_invalid_items()
            self.db_update_timestamp(timestamp)
            return Code.CUSTOM_DENYLIST_INVALIDE

        self.db_update_timestamp(timestamp)
        return Code.SUCESS

    def check_usbx_invalid_items(self):
        '''
        Check if usbx has an invalid items
        :return:
        '''

        data = self.get_usbx_invalid_denylist()
        logging.error('Invalid items from usbx {}'.format(data))
        self.db_del_by_data(data)


    def cmd_sync_denylist(self):
        '''
        Sync deny list from tbl_deny_list to U_Sandbox
        :return:
        '''

        return self._sync_denylist(self.DB_QUERY,timestamp=False)

    def cmd_apply_denylist(self,lists):
        '''
        Set list from web console to U-Sandbox
        :param list:
        :return:
        '''

        try:
            logging.debug('Load JSON file {}'.format(lists))
            json_lists = json.load(open(lists))
        except:
            logging.error('Invalid JSON format data')
            return Code.INVALID_JSON_FORMAT
        denylist = self.sort_denylist(json_lists)

        if denylist is None:
            return Code.UNSUPPORT_TYPE

        return self._sync_denylist(self.DB_INSERT, denylist)

    def cmd_modify_denylist(self,id,type,data):
        '''
        Update/Insert a denylist from webconsole
        If id is None to Insert a type, else to update an exist items
        :return:
        '''
        if data is None:
            logging.error('Please specify valid data')
            return Code.UNSUPPORT_DATA

        type = self.check_type(type)
        if type is None:
            logging.error('Please specify valid type')
            return Code.UNSUPPORT_TYPE

        if not self.librule.is_valid_all(data, type) or not is_valid_length(data):
            logging.error('Invalid items {} {}'.format(data, type))
            return Code.REASON_INVALID

        if self.is_duplicate_db(data, type):
            logging.error('Duplicate items {} {}'.format(data, type))
            return Code.REASON_DUPLICATE

        if id:
            logging.debug('Try to update an item {} {}'.format(data,id))
            return self._sync_denylist(self.DB_UPDATE, (data, id))
        else:
            logging.debug('Try to insert an item {} {}'.format(data, type))
            denylist = self.create_denylist_object(type,data)
            return self._sync_denylist(self.DB_INSERT, denylist)

    def cmd_delete_denylist(self,id):

        if id is None:
            id = ''
        return self._sync_denylist(self.DB_DELETE, id)

def set_arg():
    parser = argparse.ArgumentParser('denylist')
    parser.add_argument('--apply', action='store_true', help='Save data to database and sync to U-Sandbox')
    parser.add_argument('--modify', action='store_true', help='Enable modify or add a denylist')
    parser.add_argument('--sync', action='store_true', help='Sync denylist from db to U-Sandbox')
    parser.add_argument('--delete', action='store_true', help='Enable delete denylist')
    parser.add_argument('--dump', action='store_true', help='Dump a denylist.Use --data to specify the cnt, and default cnt is 20')
    parser.add_argument('--type', help='Denyst list type')
    parser.add_argument('--data', help='Denyst list object')
    parser.add_argument('--id', help='id or mutily id join by \',\', or empty for all denylist id')
    parser.add_argument('--list',help='JSON format denylist file')

    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(Code.OTHER_ERROR)

    return parser.parse_args()



if __name__ == '__main__':
    log_init()
    ret = None
    try:
        arg = set_arg()
        dl = DenyList()
        if arg.apply:
            ret = dl.cmd_apply_denylist(arg.list)
        elif arg.modify:
            ret = dl.cmd_modify_denylist(arg.id,arg.type,arg.data)
        elif arg.sync:
            ret = dl.cmd_sync_denylist()
        elif arg.delete:
            ret =dl.cmd_delete_denylist(arg.id)
        elif arg.dump:
            print dl.create_dumy_denylist_data(arg.data)
    except:
        logging.error(traceback.format_exc())
        ret = Code.OTHER_ERROR

    if ret:
        logging.debug('code:{}'.format(ret))
        print '{"code":%d}'%(ret)

    log_destr()
    sys.exit(ret)
